from ultralytics import YOLO
import cv2
import torch
import numpy as np
from collections import defaultdict
from frame_manager import FrameManager


class Detect:
    def __init__(self, frame_manager, model_path="E:/Innovative Practice/MyProject/YOLOv11/weights/best.pt",
                 device=None, conf_threshold=0.5):
        # 检查CUDA是否可用
        if device is None:
            device = 0 if torch.cuda.is_available() else 'cpu'
            print(f"Using device: {torch.cuda.get_device_name(0) if device == 0 else device}")
        self.model = YOLO(model_path)
        # 指定设备进行推理
        self.model.to(device)
        # 获取类别名称列表
        self.class_names = self.model.names
        self.frame_manager = frame_manager
        # 添加置信度阈值参数
        self.conf_threshold = conf_threshold

    def run_inference(self, source):
        # 设置置信度阈值
        results = self.model.predict(source, save=False, show=True, conf=self.conf_threshold)
        return results

    def process_results(self, results):
        """处理模型推理结果，提取类别名和边界框坐标"""
        detections = []

        # 遍历每个检测结果（通常图像批量处理时会有多个结果）
        for result in results:
            # 获取边界框坐标（xyxy格式）
            boxes = result.boxes.cpu().numpy()

            # 遍历当前图像中的所有检测框
            for box in boxes:
                # 获取边界框坐标
                x1, y1, x2, y2 = box.xyxy[0].astype(int)
                # 获取类别索引
                cls_idx = int(box.cls[0])
                # 获取类别名称
                cls_name = self.class_names[cls_idx]
                # 获取置信度
                confidence = float(box.conf[0])

                # 将检测信息添加到结果列表
                detections.append({
                    'class': cls_name,
                    'coordinates': ((x1 + x2) / 2, (y1 + y2) / 2),
                    'confidence': confidence
                })
                print(f"当前帧检测结果: {cls_name}, 坐标: ({(x1 + x2) / 2}, {(y1 + y2) / 2}), 置信度: {confidence:.2f}")

        return detections

    def aggregate_detections(self, all_detections):
        """聚合多帧的检测结果，计算平均坐标和最高置信度"""
        aggregated_results = defaultdict(list)

        # 聚合每类对象的检测结果
        for detections in all_detections:
            for detection in detections:
                class_name = detection['class']
                coordinates = detection['coordinates']
                confidence = detection['confidence']
                aggregated_results[class_name].append((coordinates, confidence))

        final_results = []
        for class_name, data in aggregated_results.items():
            total_confidence = 0
            avg_x = 0
            avg_y = 0
            max_confidence = 0

            for (x, y), confidence in data:
                avg_x += x
                avg_y += y
                total_confidence += confidence
                if confidence > max_confidence:
                    max_confidence = confidence

            count = len(data)
            final_results.append({
                'class': class_name,
                'coordinates': (avg_x / count, avg_y / count),
                'confidence': max_confidence
            })

        return final_results

    def detect_and_process_multiple_frames(self, num_frames=10):
        """执行推理并处理结果的便捷方法，针对多帧"""
        all_detections = []
        frame_count = 0

        while frame_count < num_frames:
            frame = self.frame_manager.get_current_frame()
            if frame is not None:
                results = self.run_inference(frame)
                detections = self.process_results(results)
                all_detections.append(detections)
                frame_count += 1
                # 显示结果
                # cv2.imshow("Camera Feed", frame)

            # if cv2.waitKey(1) & 0xFF == ord('q'):
            #     break

        # 聚合检测结果
        final_results = self.aggregate_detections(all_detections)
        return final_results


if __name__ == "__main__":
    frame_manager = FrameManager()
    # 设置置信度阈值为0.7
    detector = Detect(frame_manager, conf_threshold=0.5)
    if frame_manager.start_camera(camera_id=0):
        try:
            while True:
                results = detector.detect_and_process_multiple_frames(num_frames=10)
                print("Aggregated Detection Results for 10 frames:", results)
        finally:
            frame_manager.stop_camera()
            cv2.destroyAllWindows()



